import abc
import gher.utils.pytorch_util as ptu
from collections import OrderedDict


def _elem_or_tuple_to_variable(elem_or_tuple):
    if isinstance(elem_or_tuple, tuple):
        return tuple(
            _elem_or_tuple_to_variable(e) for e in elem_or_tuple
        )
    return ptu.from_numpy(elem_or_tuple).float()


class Trainer(object, metaclass=abc.ABCMeta):
    def __init__(self):
        self._num_train_steps = 0

    def train(self, np_batch):
        self._num_train_steps += 1
        batch = ptu.np_to_pytorch_batch(np_batch)
        self.train_from_torch(batch)

    def get_snapshot(self):
        return {}

    def get_diagnostics(self):
        return OrderedDict([
            ('num train calls', self._num_train_steps),
        ])

    @abc.abstractmethod
    def train_from_torch(self, batch):
        pass

    @abc.abstractmethod
    def end_epoch(self, epoch):
        pass
